{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "# Load environment variables from openai.env file\n",
    "load_dotenv(\"openai.env\")\n",
    "\n",
    "# Read the OPENAI_API_KEY from the environment\n",
    "api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "api_base = os.getenv(\"OPENAI_API_BASE\")\n",
    "os.environ[\"OPENAI_API_KEY\"] = api_key\n",
    "os.environ[\"OPENAI_API_BASE\"] = api_base"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 不同的Memory工具\n",
    "\n",
    "- 短时记忆：存储在内存中\n",
    "- 构建记忆实体清单\n",
    "- 接入知识图谱\n",
    "- 长对话在内存中的处理方式\n",
    "- 长时记忆实现方式\n",
    "<hr>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 内存中的短时记忆"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': 'Human: 你好，我是人类！\\nAI: 你好，我是AI，有什么可以帮助你的吗？'}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.memory import  ConversationBufferMemory\n",
    "\n",
    "memory = ConversationBufferMemory()\n",
    "memory.chat_memory.add_user_message(\"你好，我是人类！\")\n",
    "memory.chat_memory.add_ai_message(\"你好，我是AI，有什么可以帮助你的吗？\")\n",
    "\n",
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': 'Human: 你好，我是人类!\\nAI: 你好，我是AI，有什么可以帮助你的吗？\\nHuman: 我想吃鸡肉\\nAI: 好的，我帮你找找鸡肉的做法'}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#实现一个最近的对话窗口，超过窗口条数的对话将被删除\n",
    "from langchain.memory import  ConversationBufferWindowMemory\n",
    "\n",
    "memory = ConversationBufferWindowMemory(k=2)\n",
    "\n",
    "memory.save_context({\"input\":\"你好，我是人类!\"},{\"output\":\"你好，我是AI，有什么可以帮助你的吗？\"})\n",
    "memory.save_context({\"input\":\"我想吃鸡肉\"},{\"output\":\"好的，我帮你找找鸡肉的做法\"})\n",
    "\n",
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 构建记忆实体概念清单"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/tomiezhang/.pyenv/versions/3.10.12/envs/langchains/lib/python3.11/site-packages/langchain_core/_api/deprecation.py:117: LangChainDeprecationWarning: The class `langchain_community.llms.openai.OpenAI` was deprecated in langchain-community 0.0.10 and will be removed in 0.2.0. An updated version of the class exists in the langchain-openai package and should be used instead. To use it run `pip install -U langchain-openai` and import as `from langchain_openai import OpenAI`.\n",
      "  warn_deprecated(\n"
     ]
    }
   ],
   "source": [
    "from langchain.llms import  OpenAI\n",
    "from langchain.memory import ConversationEntityMemory\n",
    "\n",
    "llm = OpenAI(\n",
    "    temperature=0\n",
    ")\n",
    "\n",
    "memory = ConversationEntityMemory(llm=llm)\n",
    "_input = {\n",
    "    \"input\":\"胡八一和王胖子雪莉杨经常在一起冒险，合称盗墓铁三角.\"\n",
    "}\n",
    "memory.load_memory_variables(_input)\n",
    "memory.save_context(\n",
    "    _input,\n",
    "    {\n",
    "        \"output\":\"听起来很刺激，我也想加入他们！\"\n",
    "    }\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': 'Human: 胡八一和王胖子雪莉杨经常在一起冒险，合称盗墓铁三角.\\nAI: 听起来很刺激，我也想加入他们！',\n",
       " 'entities': {}}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({\"input\":\"铁三角是谁?\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 使用知识图谱构建记忆"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms import  OpenAI\n",
    "from langchain.memory import ConversationKGMemory\n",
    "\n",
    "llm = OpenAI(\n",
    "    temperature=0\n",
    ")   \n",
    "\n",
    "memory = ConversationKGMemory(llm=llm,return_messages=True)\n",
    "\n",
    "memory.save_context(\n",
    "    {\"input\":\"帮我找一下tomie\"},\n",
    "    {\"output\":\"对不起请问什么是tomie？\"}\n",
    ")\n",
    "\n",
    "memory.save_context(\n",
    "    {\"input\":\"tomie是一个培训讲师\"},\n",
    "    {\"output\":\"好的，我知道了。\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': [SystemMessage(content='On tomie: tomie is a training instructor.')]}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({\"input\":\"tomie是谁?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['tomie']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.get_current_entities(\"tomie最喜欢做什么事?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[KnowledgeTriple(subject='tomie', predicate='最喜欢', object_='打游戏')]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.get_knowledge_triplets(\"tomie最喜欢打游戏\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 长对话在内存中的处理方式：总结摘要以及token计算"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationSummaryMemory\n",
    "from langchain.llms import  OpenAI\n",
    "\n",
    "llm = OpenAI(\n",
    "    temperature=0\n",
    ")\n",
    "\n",
    "memory = ConversationSummaryMemory(llm=llm)\n",
    "memory.save_context(\n",
    "    {\"input\":\"帮我找一下tomie\"},\n",
    "    {\"output\":\"对不起请问什么是tomie？\"}\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\":\"tomie是一个培训讲师\"},\n",
    "    {\"output\":\"好的，我知道了。\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': '\\nThe human asks the AI to help find something called \"tomie\". The AI apologizes and asks what \"tomie\" is. The human explains that tomie is a training instructor. The AI acknowledges and says it understands.'}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nThe human asks the AI to help find information about someone named Tomie. The AI asks for clarification on who Tomie is, and the human explains that Tomie is a training instructor. The AI acknowledges this and says it understands.'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages = memory.chat_memory.messages\n",
    "#print(messages)\n",
    "memory.predict_new_summary(messages,\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "#使用ChatMessageHistory来快速获得对话摘要\n",
    "from langchain.memory import ConversationSummaryMemory\n",
    "from langchain.memory import ChatMessageHistory\n",
    "from langchain.llms import  OpenAI\n",
    "\n",
    "hisiory = ChatMessageHistory()\n",
    "hisiory.add_user_message(\"你好，我是人类！\")\n",
    "hisiory.add_ai_message(\"你好，我是AI小丸子，有什么可以帮助你的吗？\")\n",
    "\n",
    "#memory = ConversationSummaryMemory.from_messages(\n",
    "#    llm=OpenAI(temperature=0),\n",
    "#    chat_memory=hisiory,\n",
    "#    return_messages=True\n",
    "#)\n",
    "\n",
    "memory = ConversationSummaryMemory(\n",
    "    llm=OpenAI(temperature=0),\n",
    "    return_messages=True,\n",
    "    buffer=\"\\nThe AI introduces itself as AI Little Maruko and asks if there is anything it can help the human with.\",\n",
    "    chat_memory=hisiory\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': [SystemMessage(content='\\nThe AI introduces itself as AI Little Maruko and asks if there is anything it can help the human with.')]}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#memory.buffer\n",
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 当对话持续进行且对话内容很多的时候\n",
    "# 可以使用ConversationSummaryBufferMemory来存储对话摘要\n",
    "# 这是一种非常有用的方式,它会根据token的数量来自动判断是否需要进行摘要\n",
    "# 当token数量超过阈值的时候,会自动进行摘要\n",
    "# 在缓冲区中,会保留最近的k条对话\n",
    "#比较久的对话会被删除，在删除前会进行摘要"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationSummaryBufferMemory\n",
    "from langchain.llms import OpenAI\n",
    "\n",
    "llm = OpenAI(temperature=0)\n",
    "\n",
    "memory = ConversationSummaryBufferMemory(\n",
    "    llm=llm,\n",
    "    max_token_limit=10,\n",
    "    return_messages=True\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\":\"帮我找一下tomie\"},\n",
    "    {\"output\":\"对不起请问什么是tomie？\"}\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\":\"tomie是一个培训讲师\"},\n",
    "    {\"output\":\"好的，我知道了。\"}\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\":\"今天他要讲一门关于RAG的课程\"},\n",
    "    {\"output\":\"好的，我知道了。需要RAG的资料吗？\"}\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': [SystemMessage(content='\\nThe human asks the AI to help find something called \"tomie\". The AI apologizes and asks what \"tomie\" is. The human explains that tomie is a training instructor. The AI acknowledges and says it understands. The human then mentions that the instructor will be teaching a course on RAG today. The AI confirms and asks if any RAG materials are needed.')]}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conversation Token Buffer使用token长度来决定什么时候刷新内存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.memory import ConversationTokenBufferMemory\n",
    "from langchain.llms import OpenAI\n",
    "llm = OpenAI(temperature=0)\n",
    "memory = ConversationTokenBufferMemory(\n",
    "    llm=llm,\n",
    "    max_token_limit=150\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\":\"帮我找一下tomie\"},\n",
    "    {\"output\":\"对不起请问什么是tomie？\"}\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\":\"tomie是一个培训讲师\"},\n",
    "    {\"output\":\"好的，我知道了。\"}\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\":\"今天他要讲一门关于RAG的课程\"},\n",
    "    {\"output\":\"好的，我知道了。需要RAG的资料吗？\"}\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\":\"不需要资料了，谢谢\"},\n",
    "    {\"output\":\"好的，那我就不打扰你了。\"}\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': 'Human: 帮我找一下tomie\\nAI: 对不起请问什么是tomie？\\nHuman: tomie是一个培训讲师\\nAI: 好的，我知道了。\\nHuman: 今天他要讲一门关于RAG的课程\\nAI: 好的，我知道了。需要RAG的资料吗？\\nHuman: 不需要资料了，谢谢\\nAI: 好的，那我就不打扰你了。'}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "memory.load_memory_variables({})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 长时记忆的是实现方式\n",
    "\n",
    "通过向量数据库来存储之前的对话内容，有的向量数据库服务还提供自动摘要等，每次对话的时候，都会从向量数据库里查询最相关的文档或历史对话"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.vectorstores import FAISS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "memory = ConversationBufferMemory()\n",
    "memory.save_context(\n",
    "    {\"input\":\"帮我找一下tomie\"},\n",
    "    {\"output\":\"对不起请问什么是tomie？\"}\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\":\"tomie是一个培训讲师\"},\n",
    "    {\"output\":\"好的，我知道了。\"}\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\":\"今天他要讲一门关于RAG的课程\"},\n",
    "    {\"output\":\"好的，我知道了。需要RAG的资料吗？\"}\n",
    ")\n",
    "memory.save_context(\n",
    "    {\"input\":\"不需要资料了，谢谢\"},\n",
    "    {\"output\":\"好的，那我就不打扰你了。\"}\n",
    ")\n",
    "\n",
    "vectorstore = FAISS.from_texts(\n",
    "    memory.buffer.split(\"\\n\"),\n",
    "    OpenAIEmbeddings()\n",
    ")\n",
    "FAISS.save_local(vectorstore,\"test_faiss\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Document(page_content='Human: 今天他要讲一门关于RAG的课程'),\n",
       " Document(page_content='AI: 好的，我知道了。需要RAG的资料吗？'),\n",
       " Document(page_content='Human: tomie是一个培训讲师'),\n",
       " Document(page_content='AI: 好的，那我就不打扰你了。')]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "FAISS.load_local(\"test_faiss\",OpenAIEmbeddings(model=\"text-embedding-3-small\"),allow_dangerous_deserialization=True).similarity_search(\"tomie是什么职业?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'history': 'Human: tomie是一个培训讲师'}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.vectorstores import FAISS\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain.memory import VectorStoreRetrieverMemory\n",
    "r1 = FAISS.load_local(\"test_faiss\",OpenAIEmbeddings(),allow_dangerous_deserialization=True)\n",
    "r2 = r1.as_retriever(\n",
    "    search_kwargs={\"k\":1}\n",
    ")\n",
    "memory2 = VectorStoreRetrieverMemory(\n",
    "    retriever=r2\n",
    ")\n",
    "memory2.load_memory_variables({\"prompt\":\"tomie是什么职业?\"})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchains",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
     